#include "UDPSocket.h"

UDPSocket::UDPSocket(void)
	: mSocket(INVALID_SOCKET)
{
	//Initialize Critical Sections
	InitializeCriticalSection(&mCSLockSend);
	InitializeCriticalSection(&mCSLockRecv);
}

UDPSocket::UDPSocket(SOCKET s)
	: mSocket(INVALID_SOCKET)
{
	//Initialize Critical Sections
	InitializeCriticalSection(&mCSLockSend);
	InitializeCriticalSection(&mCSLockRecv);
}

UDPSocket::~UDPSocket(void)
{
	//Destroy Sockets
	if (mSocket != INVALID_SOCKET)
		closesocket(mSocket);

	//Destroy Critical Sections
	DeleteCriticalSection(&mCSLockSend);
	DeleteCriticalSection(&mCSLockRecv);
}


bool UDPSocket::RecvUDPData(char* data, int& numOfBytesRead, const std::string& ip_address, int port_num)
{
	EnterCriticalSection(&mCSLockRecv);
	
	// Receive data
	char buffer[UDP_RECV_BUFFER_SIZE];

	// Used for storing the sender's address
	sockaddr_in SenderAddr;
	socklen_t SenderAddrSize = sizeof(SenderAddr);
	
	int bytesRead = recvfrom(mSocket, buffer, UDP_RECV_BUFFER_SIZE, 0, (SOCKADDR *)&SenderAddr, &SenderAddrSize);
	
	if (bytesRead == SOCKET_ERROR || bytesRead == 0)
	{
		numOfBytesRead = 0;
		int errcode = WSAGetLastError();
		std::cout << "ERROR in recv(): " << WSAGetLastError() << std::endl;
		
		LeaveCriticalSection(&mCSLockRecv);

		return false;
	}

	std::memcpy(data, buffer, bytesRead);
	numOfBytesRead = bytesRead;

	LeaveCriticalSection(&mCSLockRecv);
	return true;
}

bool UDPSocket::SendUDPData(const char* data, int numOfBytes, const std::string& ip_address, int port_num)
{
	EnterCriticalSection(&mCSLockSend);

	// Set up the RecvAddr structure with the IP address of
	// the receiver (in this example case "192.168.1.1")
	// and the specified port number.
	sockaddr_in RecvAddr;
	RecvAddr.sin_family = AF_INET;
	RecvAddr.sin_port = ntohs(port_num);
	RecvAddr.sin_addr.s_addr = inet_addr(ip_address.c_str());
	/* Tranmsit data to get time */
	// sendto automatically binds address to the sending scket
	int bytesSent = sendto(mSocket, data, numOfBytes, 0, (SOCKADDR *)&RecvAddr, sizeof (RecvAddr));

	LeaveCriticalSection(&mCSLockSend);

	if (bytesSent)
		return true;
	else 
	{
		std::cout << "ERROR in send(): " << WSAGetLastError() << std::endl;
		return false;
	}
}

// Server-side Functionalities
bool UDPSocket::SetUpServerAddress(unsigned int &port_num)
{
	struct addrinfo hints, *result_info;

	SecureZeroMemory(&hints, sizeof(hints));

	hints.ai_family = AF_INET;			// AF_INET - Using IPv4 address family. 
	hints.ai_socktype = SOCK_DGRAM;		// Provides sequenced, reliable, two-way, connection-based byte streams
	hints.ai_protocol = IPPROTO_UDP;	// The protocol type is UDP
	hints.ai_flags = AI_PASSIVE;		// This flag is used later in getaddrinfo function.
										// It indicates that the caller intends to use the returned socket 
										// address structure in a call to the bind function.

	// Convert port to string
	char portnum_str[32];
	_itoa_s(port_num, portnum_str, 10);

	int result = getaddrinfo(NULL, portnum_str, &hints, &result_info);
	if (result)
	{
		std::cerr << "getaddrinfo failed: " << result << std::endl;
		WSACleanup();
		return false;
	}

	const unsigned int servermaxsize = 32;
	const unsigned int hostmaxsize = 1025;
	char hostname[hostmaxsize];
	char servername[servermaxsize];
	DWORD ret;

	ret = getnameinfo(result_info->ai_addr, sizeof(struct sockaddr), hostname, hostmaxsize, servername, servermaxsize, NI_NUMERICSERV);
	struct addrinfo* serverinfo2;
	getaddrinfo(hostname, NULL, &hints, &serverinfo2);

	struct sockaddr_in* addr = (struct sockaddr_in*) serverinfo2->ai_addr;
	mIPAddress = inet_ntoa(addr->sin_addr);

	freeaddrinfo(serverinfo2);

	mSocket = socket(result_info->ai_family, 
					 result_info->ai_socktype, 
					 result_info->ai_protocol);

	if (mSocket == INVALID_SOCKET)
	{
		std::cerr << "Error while trying to create socket() " << WSAGetLastError() << std::endl;
		return false;
	}

	sockaddr_in recvaddr;
	recvaddr.sin_family = AF_INET;
	recvaddr.sin_port = htons(port_num);
	recvaddr.sin_addr.s_addr = htonl(INADDR_ANY);

	int result = bind(mSocket, (SOCKADDR*)&recvaddr, sizeof(recvaddr));
	if (result == SOCKET_ERROR)
	{
		std::cerr << "Error while trying to bind socket: " << WSAGetLastError() << std::endl;
		return false;
	}

    std::cout << "Server IP Address: " << mIPAddress << std::endl;
    std::cout << "Port number: " << port_num << std::endl;

	return true;
}
